import pickle
from constant import qwen05_pretrained_model

from transformers import BertTokenizer, BertForMaskedLM, Trainer, TrainingArguments, AutoModel
import torch
from torch.utils.data.dataset import Dataset

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print("using device:", device)

model = AutoModel.from_pretrained(qwen05_pretrained_model)


# print(model)

def preprocess(data):
    # return torch.tensor(data, device=device)
    return {
        "input_ids": torch.tensor(data)

    }


class MyDatasets(Dataset):
    def __init__(self, file_path: str):
        self.data = []
        self.load_data(file_path)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return preprocess(self.data[idx])

    def load_data(self, file_path: str):
        with open(file_path, "rb") as f:
            self.data = pickle.load(f)


train_args = TrainingArguments(
    output_dir="./checkpoints",
    num_train_epochs=3,
    per_device_train_batch_size=1,
    logging_steps=500,
    save_steps=1000,
    save_total_limit=3,
    bf16=True,

)

train_dataset = MyDatasets("../preprocess_datas/train.pkl")
eval_dataset = MyDatasets("../preprocess_datas/eval.pkl")

trainer = Trainer(
    model=model,
    args=train_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
)

if __name__ == '__main__':
    trainer.train()

    trainer.save_model("./checkpoints_out")

"""
  File "C:\Users\kkk\.conda\envs\torch24\Lib\site-packages\transformers\trainer.py", line 2531, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\kkk\.conda\envs\torch24\Lib\site-packages\transformers\trainer.py", line 3675, in training_step
    loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\kkk\.conda\envs\torch24\Lib\site-packages\transformers\trainer.py", line 3752, in compute_loss
    raise ValueError(
ValueError: The model did not return a loss from the inputs, only the following keys: last_hidden_state,past_key_values. For reference, the inputs it received are input_ids.
  0%|          | 0/30000 [00:00<?, ?it/s]


"""